import streamlit as st
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from io import BytesIO

# 页面配置
st.set_page_config(
    page_title="KMeans调色板提取器",
    page_icon="🎨",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# 隐藏Streamlit默认样式
hide_streamlit_style = """
<style>
footer {visibility: hidden;}
</style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)

# 页面标题
st.markdown("<h1 style='text-align: center;'>KMeans调色板提取器</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: left;'>上传图像获取其调色板.</p>", unsafe_allow_html=True)

# 定义提取主色调的函数
def extract_colors(image, n_colors=5):
    img = image.convert('RGB')
    img = img.resize((200, 200))
    img_array = np.array(img).reshape(-1, 3)
    kmeans = KMeans(n_clusters=n_colors, random_state=42)
    kmeans.fit(img_array)
    colors = kmeans.cluster_centers_.astype(int)
    hex_colors = [f'#{r:02x}{g:02x}{b:02x}' for r, g, b in colors]
    return colors, hex_colors

# 主交互区域布局
col1, col2 = st.columns([1, 1])

with col1:
    uploaded_file = st.file_uploader("image", type=["jpg", "jpeg", "png"], label_visibility="hidden")
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, use_column_width=True)
    clear_btn, submit_btn = st.columns([1, 1])
    with clear_btn:
        st.button("Clear", key="clear")
    with submit_btn:
        submit = st.button("Submit", key="submit", type="primary")

with col2:
    palette_placeholder = st.container()
    with palette_placeholder:
        st.markdown("<h3 style='text-align: left;'>色彩十六进制码</h3>", unsafe_allow_html=True)
        color_codes = st.text_area("", "", height=150, key="color_codes", disabled=True)
        st.button("Flag", key="flag")

# 逻辑处理
if submit and uploaded_file is not None:
    image = Image.open(uploaded_file)
    colors, hex_colors = extract_colors(image)
    # 绘制调色板
    fig, ax = plt.subplots(figsize=(10, 2))
    ax.imshow([colors / 255])
    ax.axis('off')
    with col2:
        st.pyplot(fig)
    # 显示颜色码
    color_codes_text = "\n".join(hex_colors)
    with col2:
        st.text_area("", color_codes_text, height=150, key="color_codes_result", disabled=True)